import cv2
import numpy as np
import threading
from typing import Optional

class NonBlockingImageWindow:
    """
    非阻塞图像窗口（仅支持 fmt='path' 或 'rgb'）。
    - start(): 启动显示线程（__init__ 默认已启动）
    - update(img, fmt): 更新图像内容，fmt ∈ {'path','rgb'}
    - stop(timeout=None): 关闭窗口并等待显示线程退出（自动 join）
    - 自动居中：使用 tkinter 获取屏幕大小并 moveWindow 到中心
    """

    def __init__(self,
                 title: str = "Image",
                 size: Optional[tuple] = None,   # (w, h)
                 fps: int = 30,
                 auto_start: bool = True):
        self.title = title
        self.size = size            # 期望窗口大小；用于居中计算
        self.delay = max(1, int(1000 / max(1, fps)))
        self._img = None
        self._lock = threading.Lock()
        self._running = threading.Event()
        self._thread: Optional[threading.Thread] = None
        if auto_start:
            self.start()

    # ------------ 公共方法 ------------
    def start(self):
        if self._thread and self._thread.is_alive():
            return
        self._running.set()
        self._thread = threading.Thread(target=self._loop, daemon=True)
        self._thread.start()

    def update(self, img, fmt: str):
        """fmt 只能是 'path' 或 'rgb'"""
        arr = self._to_bgr(img, fmt)
        with self._lock:
            self._img = arr

    def stop(self, timeout: Optional[float] = None):
        """
        关闭窗口并等待显示线程退出（自动 join）。
        - timeout: 可选，等待秒数；None 表示一直等。
        """
        self._running.clear()
        if self._thread and threading.current_thread() is not self._thread:
            self._thread.join(timeout)
        try:
            cv2.destroyWindow(self.title)
        except Exception:
            pass

    def join(self, timeout: Optional[float] = None):
        if self._thread and threading.current_thread() is not self._thread:
            self._thread.join(timeout)

    # ------------ 内部实现 ------------
    def _loop(self):
        cv2.namedWindow(self.title, cv2.WINDOW_NORMAL)

        # 若指定了 size，则立即调整窗口大小
        if self.size:
            cv2.resizeWindow(self.title, self.size[0], self.size[1])

        # 启动时尝试居中一次（根据 self.size 或默认 640x480）
        try:
            self._center_window()
        except Exception:
            # 获取屏幕大小失败时不影响显示
            pass

        try:
            while self._running.is_set():
                frame = None
                with self._lock:
                    if self._img is not None:
                        frame = self._img

                if frame is not None:
                    cv2.imshow(self.title, frame)

                k = cv2.waitKey(self.delay)
                if k == 27:  # ESC
                    self._running.clear()
                    break
        finally:
            try:
                cv2.destroyWindow(self.title)
            except Exception:
                pass

    def _center_window(self):
        """将窗口移动到屏幕中心；依赖 tkinter 获取屏幕分辨率（失败则静默跳过）。"""
        screen_w, screen_h = self._get_screen_size_with_tk()
        if screen_w is None or screen_h is None:
            return  # 获取失败则放弃居中

        # 估计窗口大小：优先用用户给定的 size；否则使用默认 640x480
        if self.size and self.size[0] > 0 and self.size[1] > 0:
            win_w, win_h = self.size
        else:
            win_w, win_h = 640, 480

        pos_x = max(0, (screen_w - win_w) // 2)
        pos_y = max(0, (screen_h - win_h) // 2)
        cv2.moveWindow(self.title, int(pos_x), int(pos_y))

    @staticmethod
    def _get_screen_size_with_tk():
        """用 tkinter 获取屏幕分辨率；失败返回 (None, None)。"""
        try:
            import tkinter as tk
            root = tk.Tk()
            root.withdraw()
            w = root.winfo_screenwidth()
            h = root.winfo_screenheight()
            root.destroy()
            return w, h
        except Exception:
            return None, None

    def _to_bgr(self, img, fmt: str) -> np.ndarray:
        fmt = fmt.lower()
        if fmt == "path":
            arr = cv2.imread(str(img), cv2.IMREAD_UNCHANGED)
            if arr is None:
                raise ValueError(f"读取文件失败: {img}")
            return arr  # imread 已是 BGR/BGRA/Gray

        if fmt == "rgb":
            if not isinstance(img, np.ndarray):
                raise TypeError("fmt='rgb' 时 img 必须是 numpy.ndarray")
            if img.ndim != 3 or img.shape[2] != 3:
                raise ValueError("fmt='rgb' 需要形状 (H, W, 3) 的 RGB 图像")
            return cv2.cvtColor(self._normalize_to_uint8(img), cv2.COLOR_RGB2BGR)

        raise ValueError(f"不支持的 fmt: {fmt}（仅支持 'path' 和 'rgb'）")

    def _normalize_to_uint8(self, arr: np.ndarray) -> np.ndarray:
        if arr.dtype == np.uint8:
            return arr
        a = arr.astype(np.float32)
        a_min = float(np.nanmin(a))
        a_max = float(np.nanmax(a))
        if a_max == a_min:
            return np.zeros_like(a, dtype=np.uint8)
        a = (a - a_min) / (a_max - a_min) * 255.0
        return np.clip(a, 0, 255).astype(np.uint8)

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, exc_type, exc, tb):
        self.stop()
